package com.liuhengjia.springaiapply.api;

import com.liuhengjia.springaiapply.vector.Files;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Arrays;

/**
 * 加载至向量数据库 API（控制层）
 */
@Slf4j
@RestController
@RequiredArgsConstructor
public class VectorAPI {
    private final VectorStore vectorStore;

    /**
     * 加载 txt 到向量数据库
     *
     * @return code
     */
    @GetMapping("/vector/write")
    public String write() {
        InputStream resourceAsStream = getClass().getClassLoader().getResourceAsStream("who_is_xiaoyi.txt");
        StringBuilder sb = new StringBuilder();
        try (BufferedReader br = new BufferedReader(new InputStreamReader(resourceAsStream))) {
            String line;
            while ((line = br.readLine()) != null) {
                sb.append(line);
            }
        } catch (IOException e) {
            log.error(e.getMessage());
            return "500";
        }
        vectorStore.write(Arrays.stream(sb.toString().split("。")).map(Document::new).toList());
        return "200";
    }

    /**
     * 加载 pdf 到向量数据库
     *
     * @return code
     */
    @GetMapping("/vector/pdf")
    public String writePdf() {
        try {
            vectorStore.write(Files.getDocumentsFromPdf());
        } catch (Exception e) {
            log.error(e.getMessage());
            return "500";
        }
        return "200";
    }

    /**
     * 加载 markdown 到向量数据库
     *
     * @return code
     */
    @GetMapping("/vector/markdown")
    public String writeMarkdown() {
        try {
            vectorStore.write(Files.getDocumentsFormMarkdown());
        } catch (Exception e) {
            log.error(e.getMessage());
            return "500";
        }
        return "200";
    }
}
